{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "fomm-fufu.ipynb",
      "private_outputs": true,
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/eyaler/avatars4all/blob/master/fomm_fufu.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cdO_RxQZLahB"
      },
      "source": [
        "# Demo for paper \"First Order Motion Model for Image Animation\"\n",
        "\n",
        "## **Full body Tai chi and Fashion models**\n",
        "\n",
        "### Made just a little bit more accessible by Eyal Gruss (https://eyalgruss.com, eyalgruss@gmail.com)\n",
        "\n",
        "##### Original project: https://aliaksandrsiarohin.github.io/first-order-model-website\n",
        "\n",
        "##### Original notebook: https://colab.research.google.com/github/AliaksandrSiarohin/first-order-model/blob/master/demo.ipynb\n",
        "\n",
        "##### Faceswap notebook: https://colab.research.google.com/github/AliaksandrSiarohin/motion-cosegmentation/blob/master/part_swap.ipynb\n",
        "\n",
        "##### Notebook with video enhancement: https://colab.research.google.com/github/tg-bomze/Face-Image-Motion-Model/blob/master/Face_Image_Motion_Model_(Photo_2_Video)_Eng.ipynb\n",
        "\n",
        "##### Avatarify - a live vesrsion (requires local installation): https://github.com/alievk/avatarify\n",
        "\n",
        "##### Liquid Warping GAN impersonator model notebook (Impersonator) - much better for full body: https://colab.research.google.com/github/ak9250/impersonator/blob/master/impersonator.ipynb\n",
        "\n",
        "##### Liquid Warping GAN impersonator model notebook (Impersonator++) - even better: https://colab.research.google.com/drive/1bwUnj-9NnJA2EMr7eWO4I45UuBtKudg_\n",
        "\n",
        "#### **Stuff I made**:\n",
        "##### Avatars4all repository: https://github.com/eyaler/avatars4all\n",
        "##### Notebook for live webcam in the browser: https://colab.research.google.com/github/eyaler/avatars4all/blob/master/fomm_live.ipynb\n",
        "##### Notebook for talking head model: https://colab.research.google.com/github/eyaler/avatars4all/blob/master/fomm_bibi.ipynb\n",
        "##### Notebook for full body models (FOMM): https://colab.research.google.com/github/eyaler/avatars4all/blob/master/fomm_fufu.ipynb\n",
        "##### Notebook for full body models (impersonator): https://colab.research.google.com/github/eyaler/avatars4all/blob/master/ganozli.ipynb\n",
        "##### Notebook for full body models (impersonator++): https://colab.research.google.com/github/eyaler/avatars4all/blob/master/ganivut.ipynb\n",
        "##### Notebook for Wav2Lip audio based lip syncing: https://colab.research.google.com/github/eyaler/avatars4all/blob/master/melaflefon.ipynb\n",
        "##### List of more generative tools: https://j.mp/generativetools"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UCMFMJV7K-ag",
        "cellView": "form"
      },
      "source": [
        "#@title Setup\n",
        "\n",
        "%cd /content\n",
        "!git clone --depth 1 https://github.com/eyaler/first-order-model\n",
        "import os\n",
        "if not os.path.exists('/content/taichi-adv-cpk.pth.tar'):\n",
        "  !gdown https://drive.google.com/uc?id=1K_XBKgQhVY5febYaNjT18GBp9FLWMRbz -O /content/taichi-adv-cpk.pth.tar\n",
        "!wget --no-check-certificate -nc https://eyalgruss.com/fomm/taichi-adv-cpk.pth.tar -O /content/taichi-adv-cpk.pth.tar\n",
        "if not os.path.exists('/content/fashion.pth.tar'):\n",
        "  !gdown https://drive.google.com/uc?id=1RSekVOZxwAe_G48iEUd4pjg1kzHjnnUC -O /content/fashion.pth.tar\n",
        "!wget --no-check-certificate -nc https://eyalgruss.com/fomm/fashion.pth.tar -O /content/fashion.pth.tar\n",
        "!pip install -U git+https://github.com/ytdl-org/youtube-dl\n",
        "!pip install imageio==2.9.0\n",
        "!pip install imageio-ffmpeg==0.4.5\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2iGy_8v-4BrO",
        "cellView": "form"
      },
      "source": [
        "#@title Get the Driver video and Avatar image from the web\n",
        "#@markdown 1. You can change the URLs to your **own** stuff!\n",
        "#@markdown 2. Alternatively, you can upload **local** files in the next cells\n",
        "\n",
        "video_url = 'https://www.youtube.com/watch?v=NdEPmitJvaA' #@param {type:\"string\"}\n",
        "image_url = 'https://i.pinimg.com/564x/74/9e/96/749e96293399ebe1b8c99d89be522383.jpg' #@param {type:\"string\"}\n",
        "\n",
        "if video_url:\n",
        "  !rm -f /content/video.mp4\n",
        "  !youtube-dl --no-check-extensions --no-playlist -f \"bestvideo[ext=mp4][vcodec!*=av01][height<=360]+bestaudio[ext=m4a]/mp4[height<=360][vcodec!*=av01]/mp4[vcodec!*=av01]/mp4\" \"$video_url\" --merge-output-format mp4 -o /content/video\n",
        "  !mv /content/video.mp4 /content/video\n",
        "\n",
        "if image_url:\n",
        "  !wget \"$image_url\" -O /content/image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y1xCcE9XoSy4",
        "cellView": "form"
      },
      "source": [
        "#@title Optionally upload local Driver video { run: \"auto\" }\n",
        "manually_upload_video = False #@param {type:\"boolean\"}\n",
        "if manually_upload_video:\n",
        "  from google.colab import files\n",
        "  import shutil\n",
        "\n",
        "  %cd /content/sample_data\n",
        "  try:\n",
        "    uploaded = files.upload()\n",
        "  except Exception as e:\n",
        "    %cd /content\n",
        "    raise e\n",
        "\n",
        "  for fn in uploaded:\n",
        "    shutil.move('/content/sample_data/'+fn, '/content/video')\n",
        "    break\n",
        "  %cd /content"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m-zK9rlDoS7l",
        "cellView": "form"
      },
      "source": [
        "#@title Optionally upload local Avatar image { run: \"auto\" }\n",
        "manually_upload_image = False #@param {type:\"boolean\"}\n",
        "if manually_upload_image:\n",
        "  from google.colab import files\n",
        "  import shutil\n",
        "\n",
        "  %cd /content/sample_data\n",
        "  try:\n",
        "    uploaded = files.upload()\n",
        "  except Exception as e:\n",
        "    %cd /content\n",
        "    raise e\n",
        "\n",
        "  for fn in uploaded:\n",
        "    shutil.move('/content/sample_data/'+fn, '/content/image')\n",
        "    break\n",
        "  %cd /content"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uDaxVZk04PI5",
        "cellView": "form"
      },
      "source": [
        "#@title Optionally shorten Driver video\n",
        "start_seconds = 0 #@param {type:\"number\"}\n",
        "duration_seconds =  60#@param {type:\"number\"}\n",
        "start_seconds = max(start_seconds,0)\n",
        "duration_seconds = max(duration_seconds,0)\n",
        "\n",
        "if duration_seconds:\n",
        "  !mv /content/video /content/full_video\n",
        "  !ffmpeg  -ss $start_seconds -t $duration_seconds -i /content/full_video -f mp4 /content/video -y"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Oxi6-riLOgnm",
        "cellView": "form"
      },
      "source": [
        "#@title Prepare assets\n",
        "#@markdown If you ran out of RAM this means that the video is too large. You can shorten it above.\n",
        "\n",
        "center_video_to_body = True #@param {type:\"boolean\"}\n",
        "crop_video_to_body = True #@param {type:\"boolean\"}\n",
        "video_crop_expansion_factor = 1.05 #@param {type:\"number\"}\n",
        "center_image_to_body = True #@param {type:\"boolean\"}\n",
        "crop_image_to_body = False #@param {type:\"boolean\"}\n",
        "image_crop_expansion_factor = 1.05 #@param {type:\"number\"}\n",
        "video_crop_expansion_factor = max(video_crop_expansion_factor, 1)\n",
        "image_crop_expansion_factor = max(image_crop_expansion_factor, 1)\n",
        "\n",
        "import imageio\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.animation as animation\n",
        "from skimage.transform import resize\n",
        "from IPython.display import HTML, clear_output\n",
        "import cv2\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "hog = cv2.HOGDescriptor()\n",
        "hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())\n",
        "\n",
        "def fix_dims(im):\n",
        "    if im.ndim == 2:\n",
        "        im = np.tile(im[..., None], [1, 1, 3])\n",
        "    return im[...,:3]\n",
        "\n",
        "def get_crop(im, center_body=True, crop_body=True, expansion_factor=1, rects=None):\n",
        "    im = fix_dims(im)\n",
        "    if (center_body or crop_body) and rects is None:\n",
        "        rects, _ = hog.detectMultiScale(im, winStride=(4, 4),padding=(8,8), scale=expansion_factor)\n",
        "    if (center_body or crop_body) and rects is not None and len(rects):\n",
        "        x0,y0,w,h = sorted(rects, key=lambda x: x[2]*x[3])[-1]\n",
        "        if crop_body:\n",
        "            x0 += w//2-h//2\n",
        "            x1 = x0+h\n",
        "            y1 = y0+h\n",
        "        else:\n",
        "            img_h,img_w = im.shape[:2]\n",
        "            x0 += (w-img_h)//2\n",
        "            x1 = x0+img_h\n",
        "            y0 = 0\n",
        "            y1 = img_h\n",
        "    else:\n",
        "        h,w = im.shape[:2]\n",
        "        x0 = (w-h)//2\n",
        "        x1 = (w+h)//2\n",
        "        y0 = 0\n",
        "        y1 = h\n",
        "    return int(x0),int(x1),int(y0),int(y1)\n",
        "\n",
        "def pad_crop_resize(im, x0=None, x1=None, y0=None, y1=None, new_h=256, new_w=256):\n",
        "    im = fix_dims(im)\n",
        "    h,w = im.shape[:2]\n",
        "    if x0 is None:\n",
        "      x0 = 0\n",
        "    if x1 is None:\n",
        "      x1 = w\n",
        "    if y0 is None:\n",
        "      y0 = 0\n",
        "    if y1 is None:\n",
        "      y1 = h\n",
        "    if x0<0 or x1>w or y0<0 or y1>h:\n",
        "        im = np.pad(im, pad_width=[(max(-y0,0),max(y1-h,0)),(max(-x0,0),max(x1-w,0)),(0,0)], mode='edge')\n",
        "    return resize(im[max(y0,0):y1-min(y0,0),max(x0,0):x1-min(x0,0)], (new_h, new_w))\n",
        "\n",
        "\n",
        "source_image = imageio.imread('/content/image')\n",
        "source_image = pad_crop_resize(source_image, *get_crop(source_image, center_body=center_image_to_body, crop_body=crop_image_to_body, expansion_factor=image_crop_expansion_factor))\n",
        "\n",
        "with imageio.get_reader('/content/video', format='mp4') as reader:\n",
        "  fps = reader.get_meta_data()['fps']\n",
        "\n",
        "  driving_video = []\n",
        "  rects = None\n",
        "  try:\n",
        "      for i,im in enumerate(reader):\n",
        "          if not crop_video_to_body:\n",
        "              break\n",
        "          rects, _ = hog.detectMultiScale(im, winStride=(4, 4),padding=(8,8), scale=video_crop_expansion_factor)\n",
        "          if rects is not None and len(rects):\n",
        "              break\n",
        "      x0,x1,y0,y1 = get_crop(im, center_body=center_video_to_body, crop_body=crop_video_to_body, expansion_factor=video_crop_expansion_factor, rects=rects)\n",
        "      reader.set_image_index(0)\n",
        "      for im in reader:\n",
        "          driving_video.append(pad_crop_resize(im,x0,x1,y0,y1))\n",
        "  except RuntimeError:\n",
        "      pass\n",
        "\n",
        "def vid_display(source, driving, generated=None):\n",
        "    fig = plt.figure(figsize=(8 + 4 * (generated is not None), 6))\n",
        "\n",
        "    ims = []\n",
        "    for i in range(len(driving)):\n",
        "        cols = [source]\n",
        "        cols.append(driving[i])\n",
        "        if generated is not None:\n",
        "            cols.append(generated[i])\n",
        "        im = plt.imshow(np.concatenate(cols, axis=1), animated=True)\n",
        "        plt.axis('off')\n",
        "        ims.append([im])\n",
        "\n",
        "    ani = animation.ArtistAnimation(fig, ims, interval=50, repeat_delay=1000)\n",
        "    plt.close()\n",
        "    return ani\n",
        "\n",
        "clear_output()\n",
        "if rects is not None and len(rects):\n",
        "    print('first found body in frame %d'%i)\n",
        "HTML(vid_display(source_image, driving_video).to_html5_video())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SB12II11kF4c",
        "cellView": "form"
      },
      "source": [
        "#@title Animate\n",
        "\n",
        "choose_model = 'Tai chi' #@param ['Tai chi', 'Fashion'] {type:\"raw\"}\n",
        "exaggerate_factor = 1 #@param {type:\"slider\", min:0.1, max:5, step:0.1}\n",
        "adapt_movement_scale = True #@param {type:\"boolean\"}\n",
        "use_relative_movement = True #@param {type:\"boolean\"}\n",
        "use_relative_jacobian = True #@param {type:\"boolean\"}\n",
        "\n",
        "%cd /content/first-order-model\n",
        "import torch\n",
        "from tqdm import tqdm\n",
        "from scipy.spatial import ConvexHull\n",
        "from demo import load_checkpoints\n",
        "from skimage import img_as_ubyte\n",
        "\n",
        "def full_normalize_kp(kp_source, kp_driving, kp_driving_initial, adapt_movement_scale=False,\n",
        "                 use_relative_movement=False, use_relative_jacobian=False, exaggerate_factor=1):\n",
        "    if adapt_movement_scale:\n",
        "        source_area = ConvexHull(kp_source['value'][0].data.cpu().numpy()).volume\n",
        "        driving_area = ConvexHull(kp_driving_initial['value'][0].data.cpu().numpy()).volume\n",
        "        adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)\n",
        "    else:\n",
        "        adapt_movement_scale = 1\n",
        "\n",
        "    kp_new = {k: v for k, v in kp_driving.items()}\n",
        "\n",
        "    if use_relative_movement:\n",
        "        kp_value_diff = (kp_driving['value'] - kp_driving_initial['value'])\n",
        "        kp_value_diff *= adapt_movement_scale * exaggerate_factor\n",
        "        kp_new['value'] = kp_value_diff + kp_source['value']\n",
        "\n",
        "        if use_relative_jacobian:\n",
        "            jacobian_diff = torch.matmul(kp_driving['jacobian'], torch.inverse(kp_driving_initial['jacobian']))\n",
        "            kp_new['jacobian'] = torch.matmul(jacobian_diff, kp_source['jacobian'])\n",
        "\n",
        "    return kp_new\n",
        "\n",
        "def make_animation(source_image, driving_video, generator, kp_detector, adapt_movement_scale=False,\n",
        "                 use_relative_movement=False, use_relative_jacobian=False, cpu=False, exaggerate_factor=1):\n",
        "    with torch.no_grad():\n",
        "        predictions = []\n",
        "        source = torch.tensor(source_image[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2)\n",
        "        if not cpu:\n",
        "            source = source.cuda()\n",
        "        driving = torch.tensor(np.array(driving_video)[np.newaxis].astype(np.float32)).permute(0, 4, 1, 2, 3)\n",
        "        kp_source = kp_detector(source)\n",
        "        kp_driving_initial = kp_detector(driving[:, :, 0])\n",
        "\n",
        "        for frame_idx in tqdm(range(driving.shape[2])):\n",
        "            driving_frame = driving[:, :, frame_idx]\n",
        "            if not cpu:\n",
        "                driving_frame = driving_frame.cuda()\n",
        "            kp_driving = kp_detector(driving_frame)\n",
        "            kp_norm = full_normalize_kp(kp_source=kp_source, kp_driving=kp_driving,\n",
        "                                   kp_driving_initial=kp_driving_initial, adapt_movement_scale=adapt_movement_scale, use_relative_movement=use_relative_movement,\n",
        "                                   use_relative_jacobian=use_relative_jacobian, exaggerate_factor=exaggerate_factor)\n",
        "            out = generator(source, kp_source=kp_source, kp_driving=kp_norm)\n",
        "\n",
        "            predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])\n",
        "    return predictions\n",
        "\n",
        "if choose_model == 'Tai chi':\n",
        "  generator, kp_detector = load_checkpoints(config_path='/content/first-order-model/config/taichi-adv-256.yaml', checkpoint_path='/content/taichi-adv-cpk.pth.tar')\n",
        "else:\n",
        "  generator, kp_detector = load_checkpoints(config_path='/content/first-order-model/config/fashion-256.yaml', checkpoint_path='/content/fashion.pth.tar')\n",
        "predictions = make_animation(source_image, driving_video, generator, kp_detector, adapt_movement_scale=adapt_movement_scale, use_relative_movement=use_relative_movement,\n",
        "                                   use_relative_jacobian=use_relative_jacobian, exaggerate_factor=exaggerate_factor)\n",
        "\n",
        "imageio.mimsave('/content/generated.mp4', [img_as_ubyte(frame) for frame in predictions], fps=fps)\n",
        "!ffmpeg -i /content/generated.mp4 -i /content/video -c:v libx264 -c:a aac -map 0:v -map 1:a? -pix_fmt yuv420p -profile:v baseline -movflags +faststart /content/final.mp4 -y\n",
        "#video can be downloaded from /content/final.mp4\n",
        "\n",
        "clear_output()\n",
        "HTML(vid_display(source_image, driving_video, predictions).to_html5_video())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lH4oJcBOKZcf",
        "cellView": "form"
      },
      "source": [
        "#@title Download\n",
        "#@markdown 1. If it fails try running this cell again.\n",
        "#@markdown 2. Alternatively, you can manually download \"final.mp4\" from the folder on the left (click \"Refresh\" if missing).\n",
        "\n",
        "print() #see https://github.com/googlecolab/colabtools/issues/468\n",
        "from google.colab import files\n",
        "files.download('/content/final.mp4') #fails for Firefox private window"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}
